package com.stevin.tool.inject;

import javassist.*;

import java.lang.instrument.ClassFileTransformer;
import java.security.ProtectionDomain;

import static com.stevin.tool.utils.ClassPermissionUtils.*;

/**
 * @author mars
 * @date 2018/11/03
 */
public class CodeInterceptor implements ClassFileTransformer {

    private static String DEFAULT_POINT_CLASSNAME = "*";

    public CodeInterceptor(String agentArgs) {
        DEFAULT_POINT_CLASSNAME = agentArgs;
    }

    @Override
    public byte[] transform(ClassLoader loader,
                            String className,
                            Class<?> classBeingRedefined,
                            ProtectionDomain protectionDomain,
                            byte[] classfileBuffer) {
        try {
            if (isThirdPartyClass(className)) {
                return new byte[0];
            }
            if (!(className.startsWith(DEFAULT_POINT_CLASSNAME) || DEFAULT_POINT_CLASSNAME.equalsIgnoreCase("*"))) {
                return new byte[0];
            }


            ClassPool classPool = ClassPool.getDefault();

            String conClass = className.replace("/", ".");
            CtClass ctClass;

            try {
                ctClass = classPool.getCtClass(conClass);
            } catch (Throwable e) {
                return new byte[0];
            }

            if (isFrozen(ctClass.getModifiers())) {
                return new byte[0];
            }

            if (ctClass.isInterface() || ctClass.isAnnotation()) {
                return new byte[0];
            }

            CtConstructor[] constructors = ctClass.getDeclaredConstructors();

            for (CtConstructor ctConstructor : constructors) {

                if (ctConstructor.callsSuper()) {
                    try {
                        ctConstructor.insertBefore("{System.out.println(\" 开始  \"+Thread.currentThread().getName()+\" : \"+ java.lang.invoke.MethodHandles.lookup().lookupClass() +\"." + ctConstructor.getName() + "()\");}");
                        ctConstructor.insertAfter("{System.out.println(\" 结束  \"+Thread.currentThread().getName()+\" : \"+ java.lang.invoke.MethodHandles.lookup().lookupClass() +\"." + ctConstructor.getName() + "()\");}");
                    } catch (Throwable e) {
                        e.printStackTrace();
                    }
                    continue;
                }
                String methodParam = getBehaviorParam(ctConstructor);

                try {
                    ctConstructor.insertBefore("{System.out.println(\" 开始  \"+Thread.currentThread().getName()+\" : \"+ java.lang.invoke.MethodHandles.lookup().lookupClass() +\"." + ctConstructor.getName() + "(" + methodParam + ")\");}");
                    ctConstructor.insertAfter("{System.out.println(\" 结束  \"+Thread.currentThread().getName()+\" : \"+ java.lang.invoke.MethodHandles.lookup().lookupClass() +\"." + ctConstructor.getName() + "\");}");
                } catch (Throwable e) {
                    e.printStackTrace();
                }
            }

            CtMethod[] methods = ctClass.getDeclaredMethods();

            for (CtMethod method : methods) {
                if (isFrozen(method.getModifiers())) {
                    continue;
                }

                if (isNativeMethod(method.getName())) {
                    continue;
                }

                String methodParam = getBehaviorParam(method);

                String methodReturn = getMethodReturn(method);

                try {
                    method.addLocalVariable("start", CtClass.longType);

                    method.insertBefore("{  start=System.currentTimeMillis();" +
                            " System.out.println(\" 开始  \"+Thread.currentThread().getName()+\" : " + getCurrentSimpleEndpoint(method.getName(), methodParam) + "\");}");
                    method.insertAfter("{ try { " +
                            "System.out.println(\" 结束  \"+Thread.currentThread().getName()+\" : " + getCurrentSimpleEndpoint(method.getName(), methodParam) + " 返回值:\"+ $_ +\" 执行耗时: \"+(System.currentTimeMillis()-start)+\"\"); } " +
                            "catch (java.lang.NullPointerException e) " +
                            "{" +
                            " System.out.println(\" 结束  \"+Thread.currentThread().getName()+\" : " + getCurrentSimpleEndpoint(method.getName(), methodParam) + " 返回类型:" + methodReturn + "\"); " +
                            "}" +
                            "}");
                } catch (Throwable e) {
                    e.printStackTrace();
                }
            }

            return ctClass.toBytecode();
        } catch (Exception e) {
            e.printStackTrace();
        }
        return new byte[0];
    }

    private static Object get() {
        try {
            return "$_";
        } catch (Throwable e) {
            return "";
        }
    }


    private String getMethodReturn(CtMethod method) {
        try {
            return method.getReturnType().getName();
        } catch (Exception e) {

        }
        return "";
    }

    private static String getCurrentSimpleEndpoint(String methodName, String methodParam) {
        return "\"+java.lang.invoke.MethodHandles.lookup().lookupClass()+\"" + "." + methodName + "(" + methodParam + ")";
    }


    private static String getBehaviorParam(CtBehavior method) {
        StringBuilder paramStr = new StringBuilder("");
        try {
            if (method.getParameterTypes().length == 0) {
                return "";
            }
            int count = 1;
            for (CtClass paramType : method.getParameterTypes()) {
                paramStr.append(paramType.getSimpleName());
                paramStr.append(" \"+ $" + count + " +\" , ");
                count++;
            }
            if (paramStr.length() > 2) {
                return paramStr.substring(0, paramStr.length() - 2);
            }
        } catch (Exception e) {
            return "";
        }
        return paramStr.toString();
    }

}
